
library(policyPlot)
library(dplyr)
library(plyr)
library(data.table)
library(RColorBrewer)

#dfa_colors <- c("#7570B3", "#1B9E77", "#386CB0", "#E7298A", "#D95F02", "#E6AB02", "pink", "cornflowerblue", "mistyrose3", "seagreen1", "slategray")
dfa_colors <- c("#000000", "#9A9A9A", "#4D4D4D", "#D7D7D7", "#737373", "lightgrey", "#8E8E8E", "#C8C8C8", "#E6E6E6","#C4C1C2", "#F0F0F0")

dfa_shares_stacked <- function(plotdata, dots = FALSE, scf_dots = scf_qtrs, Col = dfa_colors, legend = TRUE, legend.col = Col,
                               legend.x.loc = "top", legend.y.loc = NULL, legend.text = NULL, legend.fill = dfa_colors, 
                               legend.text.col = "black", legend.ncol = ceiling(ngroups/2), legend.cex = .6125, legend.x.intersp = .5,
                               Enddatelab = FALSE, Y2lab = "Percent of aggregate", Zeroline = TRUE, Y2lim = c(-10,115),
                               Y2at = seq(0,100,20), Slab_ends = FALSE, Slab_skip = c(), ...) {
  
  # Prepare series labels (if Slab_ends == TRUE)
  
  # Slab_skip <- c()
  # Slab_skip <- c(1)
  # plotdata <- nw_shares_bynw
  
  len <- length(plotdata)
  ngr <- len - 1 - length(Slab_skip)
  mark <- seq(1,(len-1))
  
  if (!is.null(Slab_skip)) {
    mark <- mark[-Slab_skip]
  } 
  
  if (Slab_ends == TRUE) {
    #EMD added data.table 10/21
    Slabvec <- c(paste0(as.numeric(round(data.table::first(plotdata[2:len])))[mark], "%"), paste0(as.numeric(round(data.table::last(plotdata[2:len])))[mark], "%"))
    Slabloc.y <- c((cumsum(as.numeric(data.table::first(plotdata[2:len])))[mark] + c(0,lag(cumsum(as.numeric(data.table::first(plotdata[2:len]))))[2:(len-1)])[mark])/2,
                   (cumsum(as.numeric(data.table::last(plotdata[2:len])))[mark] + c(0,lag(cumsum(as.numeric(data.table::last(plotdata[2:len]))))[2:(len-1)])[mark])/2)
    
    Slabloc.x <- c(rep(decimal_date(data.table::first(nw_shares_bynw$date)) + .75,ngr), rep(decimal_date(data.table::last(nw_shares_bynw$date)) - 3,ngr))
    Slabadj <- c(0,.5)
    
  } else {
    Slabvec <- NULL
    Slabloc.y <- NULL
    Slabloc.x <- NULL
    Slabadj <- NULL
    
  }
  
  pinfo <- rplot.line(
    plotdata = plotdata,
    Col = Col,
    stack.line = TRUE,
    legend = FALSE,
    Enddatelab = Enddatelab,
    Y2lab = Y2lab,
    Zeroline = Zeroline,
    Y2lim = Y2lim,
    Y2at = Y2at,
    Slabvec = Slabvec,
    Slabloc.y = Slabloc.y,
    Slabloc.x = Slabloc.x,
    Slabadj = Slabadj,
    ...
  )
  
  ngroups <- ncol(pinfo$plotdata.right) -1
  
  
  if(dots == TRUE) {
    points(x = rep(round_any(decimal_date(as.Date(scf_dots, frac=1)), .25), ngroups), y = as.numeric(as.matrix(pinfo$plotdata.right %>% filter(date %in% as.Date(scf_dots, frac=1)) %>% select(-date))), pch = 20, cex = .5)
  }
  
  if (legend == TRUE & dots == FALSE) {
    legend(x = legend.x.loc, y = legend.y.loc, 
           legend = legend.text,
           fill = legend.col,
           text.col = legend.text.col,
           ncol = legend.ncol,
           cex = legend.cex,
           bty = "n",
           border = FALSE,
           x.intersp = legend.x.intersp)
    
  }
  
  if (legend == TRUE & dots == TRUE) {
    legend(x = legend.x.loc, y = legend.y.loc, 
           legend = c(legend.text, "SCF"),
           fill = c(legend.col, NA),
           text.col = legend.text.col,
           ncol = legend.ncol,
           cex = legend.cex,
           bty = "n",
           border = FALSE,
           x.intersp = legend.x.intersp,
           pch = c(rep(NA,ngroups), 20),
           col = c(rep(NA,ngroups), "black"),
           merge = FALSE,
           inset = 0,
           xjust = 0,
           yjust = 1)
    
  }
  
  
}

dfa_levels_stacked <- function(plotdata, dots = FALSE, scf_dots = scf_qtrs, Col = dfa_colors, legend = TRUE, legend.col= Col,
                               legend.x.loc = "topleft", legend.y.loc = NULL, legend.text = NULL, legend.fill = dfa_colors, 
                               legend.text.col = "black", legend.ncol = 2, legend.cex = .6125, legend.x.intersp = .5, legend.inset = .05,
                               Enddatelab = FALSE, Y2lab = "Trillions of nominal dollars", Zeroline = TRUE, ...) {
  
  
  pinfo <- rplot.line(
    plotdata = plotdata,
    Col = Col,
    stack.line = TRUE,
    legend = FALSE,
    Enddatelab = Enddatelab,
    Y2lab = Y2lab,
    Zeroline = Zeroline,
    ...
  )
  
  ngroups <- ncol(pinfo$plotdata.right) -1
  
  if(dots == TRUE) {
    points(x = rep(round_any(decimal_date(as.Date(scf_dots, frac=1)), .25), ngroups), y = as.numeric(as.matrix(pinfo$plotdata.right %>% filter(date %in% as.Date(scf_dots, frac=1)) %>% select(-date))), pch = 20, cex = .5)
  }
  
  if (legend == TRUE & dots == FALSE) {
    legend(x = legend.x.loc, y = legend.y.loc, 
           legend = legend.text,
           col = legend.col,
           lwd = 1.5,
           text.col = legend.text.col,
           cex = legend.cex,
           bty = "n",
           border = FALSE,
           inset = legend.inset)
    
  }
  
  if (legend == TRUE & dots == TRUE) {
    legend(x = legend.x.loc, y = legend.y.loc, 
           legend = c(legend.text, "SCF"),
           col = c(legend.col, "black"),
           lwd = c(rep(1.5, ngroups), NA),
           text.col = legend.text.col,
           cex = legend.cex,
           pch = c(rep(NA,ngroups), 20),
           bty = "n",
           border = FALSE,
           merge = FALSE,
           inset = legend.inset)
    
  }
  
  
}

balancesheeter <- function(df, date) {
  date_input = date
  df_balance <- df %>% mutate(morts = -morts, other_liab = -liab - morts, other_fin = fin - corp_equ_mfs - non_corp_eq, non_fin = non_fin, bus_eq = corp_equ_mfs + non_corp_eq) %>%
    select(c(date,cat,bus_eq,other_fin,non_fin,morts,other_liab,networth)) %>%
    group_by(cat) %>%
    gather(key = class, value = value, -c(date,cat)) %>%
    arrange(date) %>%
    subset(date == date_input)
  df_balance$class <- factor(df_balance$class, levels = c("bus_eq", "other_fin", "non_fin","morts","other_liab", "networth") )
  # levels(df_balance$class) <- list("Business Equity" = "bus_eq", "Other Fin Assets" = "other_fin","NonFin Assets" = "non_fin",  "Mortgages" = "morts","Other Liabilities" = "other_liab",  "Net Worth" = "networth")
  return(df_balance)
}

dfa_balance_sheet <- function(plotdata, 
                              legend.text, 
                              xaxis.label, 
                              networth.dots = TRUE, 
                              networth.pch = 20,
                              Col = c("#000000", "#7D7D7D", "#9D9D9D", "#4D4D4D", "#BDBDBD"),#c("steelblue4", "lightskyblue2", "palegreen3", "red3" ,"indianred1"), 
                              cex.yaxis = .75,
                              width = 1, 
                              space = .15, 
                              upper.ext = 1.2, 
                              lower.ext = 1.1,
                              Ylab = "Trillions of nominal dollars", 
                              Ylab.cex = .6,
                              Ylab.line = .1,
                              legend = TRUE,
                              legend.ncol = 2,
                              legend.cex = .5,
                              xaxis.cex = .65,
                              xaxis.padj = .5,
                              Zeroline = TRUE,
                              leftgap =0,
                              rightgap =0,
                              ylim.special = NULL,
                              ...) {
  
  
  if (networth.dots == TRUE) {
    networth <- plotdata[which(plotdata$class == "networth"),]
    plotdata <- plotdata[which(plotdata$class != "networth"),]
  }
  
  bars <- data.matrix(plotdata[,2:length(plotdata)])
  rownames(bars) <- plotdata[["class"]]
  nbars <- ncol(bars)
  nclass <- nrow(bars)
  
  # Need to plot positive and negative portions separately
  pos <- bars
  neg <- bars
  
  pos[pos<0] <- 0
  neg[neg>0] <- 0
  myrange <- c(min(colSums(neg)),max(colSums(pos)))
  
  
  # Initialize the info lists and set the limits
  xinfo = NULL
  yinfo = NULL
  xinfo$xlim <- c(-space - leftgap, nbars * width + nbars * space + 2 * space + rightgap)
  yinfo$ylim = c(myrange[1] * lower.ext, myrange[2] * upper.ext)
  
  if (!is.null(ylim.special)) {
    yinfo$ylim <- ylim.special
  }
  
  # Make the barplot
  # Plot positive portion
  pinfo <- barplot(height = pos,
                   col = Col,
                   ylim = yinfo$ylim,
                   xlim = xinfo$xlim,
                   axisnames = F,
                   axes = F,
                   width = width,
                   space = space,
                   beside = F, 
                   cex.axis = cex.yaxis,
                   names.arg = rep("", nbars),
                   legend.text = FALSE,
                   ...)
  # Plot negative portion
  barplot(height = neg, add = TRUE, xlim = xinfo$xlim, ylim = rev(yinfo$ylim), beside = F,
          col = Col,
          axes = FALSE,
          names.arg = rep("", nbars),
          width = width,
          space = space)
  
  if (networth.dots == TRUE) {
    points(x = pinfo, y = as.numeric(networth[,2:(nbars + 1)]), pch = networth.pch)
  }
  
  mtext(Ylab, side = 3, line = Ylab.line, adj = 1, cex = Ylab.cex, at = xinfo$xlim[2])
  
  if (legend == TRUE) {
    
    if (networth.dots == TRUE) {
      legend(x="top", legend = legend.text,
             fill = c(Col, NA),
             border = NA,
             pch = c(rep(NA,nclass), networth.pch),
             bty = "n", ncol = legend.ncol, cex = legend.cex, merge = FALSE)
    }
    
    if (networth.dots == FALSE) {
      legend(x="top", legend = legend.text,
             fill = Col,
             border = NA,
             bty = "n", ncol = legend.ncol, cex = legend.cex)
    }
  }
  
  # Label x axis 
  axis_label_loc <- pinfo
  axis_label <- xaxis.label
  axis(1, axis_label_loc, axis_label, cex.axis = xaxis.cex, tck = 0, padj = xaxis.padj)
  # padj = -2.1
  # Make the axis
  policyPlot::makeyaxis(xinfo, yinfo, side = 2, noticklabs = T)
  policyPlot::makeyaxis(xinfo, yinfo, side = 4, noticklabs = F)
  policyPlot::makexaxis(xinfo, yinfo, Majortics = FALSE, Minortics = FALSE)
  
  # Add zero line
  if (Zeroline == TRUE) {
    clip(xinfo$xlim[1],xinfo$xlim[2],-1,1)
    # clip(-space, width * nbars + space * (nbars+1),-1,1)
    abline(h=0, xpd = F)
  }
  
}